'''
Author: DiChen
Date: 2024-08-22 22:28:36
LastEditors: DiChen
LastEditTime: 2024-08-23 21:22:26
'''
"""
Author: DiChen
Date: 2024-08-20 15:55:46
LastEditors: DiChen
LastEditTime: 2024-08-21 22:55:06
"""

import requests
from rembg import remove
import argparse
import sys
from io import BytesIO
import random
import string
import json


def remove_background(url):
    response = requests.get(url)
    input = response.content
    output = remove(input)
    # 将处理后的内容包装为一个内存中的文件对象
    file_object = BytesIO(output)
    file_object.name = (
        generate_random_string(10) + ".png"
    )  # 设置文件名（可根据实际需求命名）

    # 将文件以multipart/form-data的形式POST出去
    post_url = "http://112.4.132.6:8083/data"  # 替换为实际的目标URL
    files = {"datafile": (file_object.name, file_object)}

    res = requests.post(post_url, files=files)
    if res.status_code == 200:
        text = json.loads(res.text)
        url = (
            "http://112.4.132.6:8083/data/"
            + text["data"]["id"]
        )
        print("success!")
        print("url:" + url)
        return url
    else:
        print("fail!")
        sys.exit(1)


def generate_random_string(length):
    # 生成包含字母和数字的随机字符串
    characters = string.ascii_letters + string.digits
    random_string = "".join(random.choices(characters, k=length))
    return random_string


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Remove background from an image")
    parser.add_argument("url", type=str, help="URL of the image")
    args = parser.parse_args()
    print("initializing...")
    remove_background(args.url)